import sqlite3
from datetime import datetime
import tushare as ts

from dto.strategy_stock_day import StrategyStockDay
from .database_pool import DatabasePool
from peewee import *
from dto.stock_trade_info_dto import StockTradeInfoDTO


class StockTradeInfo(Model):
    # 股票代码
    stock_code = CharField()
    # 交易日期
    trade_date = DateField()
    # 开盘价
    open = FloatField()
    # 最高价
    high = FloatField()
    # 最低价
    low = FloatField()
    # 收盘价
    close = FloatField()
    # 前收盘价
    pre_close = FloatField()
    # 涨跌额
    change = FloatField()
    # 涨跌幅
    pct_chg = FloatField()
    # 成交量
    vol = FloatField()
    # 成交金额
    amount = FloatField()

    class Meta:
        database = DatabasePool(max_connections=5).database  # 指定数据库连接
        table_name = "ant_monitor_day"  # 指定表名


class StockDateDayDas:
    # 增加数据
    def add_stock(self, df) -> StockTradeInfoDTO:
        pool = DatabasePool()
        conn = pool.connect()
        cursor = conn.cursor()

        df_sorted = df.sort_values(by="trade_date", ascending=True).reset_index(
            drop=True
        )
        AntLogger.infodf_sorted)
        trade_infos = []
        for index, row in df_sorted.iterrows():
            stockInfo = StockTradeInfo(
                stock_code=row["ts_code"],
                trade_date=row["trade_date"],
                open=row["open"],
                high=row["high"],
                low=row["low"],
                close=row["close"],
                pre_close=row["pre_close"],
                change=row["change"],
                pct_chg=row["pct_chg"],
                vol=row["vol"],
                amount=row["amount"],
            )
            trade_infos.append(stockInfo)

        StockTradeInfo.bulk_create(trade_infos)
        return trade_infos[0]

    def delete_bycode(self, stack_code):
        delete_query = StockTradeInfo.delete().where(
            StockTradeInfo.stock_code == stack_code
        )
        delete_query.execute()  # 执行删除操作

    def search_by_day(self, stock_code, day_num):
        today = datetime.now().date()
        stock_trade_info_list = (
            StockTradeInfo.select()
            .where(
                (StockTradeInfo.stock_code == stock_code)
                & (StockTradeInfo.trade_date.between(day_num, today))
            )
            .order_by(StockTradeInfo.trade_date.desc())
        )

        return [
            self.model_to_dto(stock_trade_info)
            for stock_trade_info in stock_trade_info_list
        ]

    def search_by_day_all(self, stock_code):
        today = datetime.now().date()
        stock_trade_info_list = (
            StockTradeInfo.select()
            .where(StockTradeInfo.stock_code == stock_code)
            .order_by(StockTradeInfo.trade_date.asc())
        )

        return [
            self.model_to_dto(stock_trade_info)
            for stock_trade_info in stock_trade_info_list
        ]

    def get_last_trade_date(self, stock_code):
        """
        查询某个 stock_code 在 StockDateDayDas 中的最后一条交易日期
        """
        last_record = (
            StockTradeInfo.select()
            .where(StockTradeInfo.stock_code == stock_code)
            .order_by(StockTradeInfo.trade_date.desc())
            .first()
        )
        return last_record.trade_date if last_record else None

    def model_to_dto(self, stock_trade_info: StockTradeInfo):
        return StrategyStockDay(
            stock_code=stock_trade_info.stock_code,
            trade_date=stock_trade_info.trade_date,
            open=stock_trade_info.open,
            high=stock_trade_info.high,
            low=stock_trade_info.low,
            close=stock_trade_info.close,
            pre_close=stock_trade_info.pre_close,
            change=stock_trade_info.change,
            pct_chg=stock_trade_info.pct_chg,
            vol=stock_trade_info.vol,
            amount=stock_trade_info.amount,
        )

    def dto_to_model(self, dto: StrategyStockDay):
        return StockTradeInfo.create(
            stock_code=dto.stock_code,
            trade_date=dto.trade_date,
            open=dto.open,
            high=dto.high,
            low=dto.low,
            close=dto.close,
            pre_close=dto.pre_close,
            change=dto.change,
            pct_chg=dto.pct_chg,
            vol=dto.vol,
            amount=dto.amount,
        )
